package com.smartaccounting.dao;

import com.smartaccounting.model.Budget;
import com.smartaccounting.util.DatabaseUtil;

import java.math.BigDecimal;
import java.sql.*;
import java.util.ArrayList;
import java.util.List;

public class BudgetDAO {

    // 根据ID获取预算
    public Budget getBudgetById(int budgetId) throws SQLException {
        String sql = "SELECT * FROM budgets WHERE budget_id = ?";
        try (Connection conn = DatabaseUtil.getConnection();
             PreparedStatement stmt = conn.prepareStatement(sql)) {
            stmt.setInt(1, budgetId);
            ResultSet rs = stmt.executeQuery();
            if (rs.next()) {
                return extractBudgetFromResultSet(rs);
            }
        }
        return null;
    }

    // 获取用户的所有预算
    public List<Budget> getBudgetsByUserId(int userId) throws SQLException {
        List<Budget> budgets = new ArrayList<>();
        String sql = "SELECT * FROM budgets WHERE user_id = ? ORDER BY month_year DESC, category";
        try (Connection conn = DatabaseUtil.getConnection();
             PreparedStatement stmt = conn.prepareStatement(sql)) {
            stmt.setInt(1, userId);
            ResultSet rs = stmt.executeQuery();
            while (rs.next()) {
                budgets.add(extractBudgetFromResultSet(rs));
            }
        }
        return budgets;
    }

    // 获取用户指定月份的预算
    public List<Budget> getBudgetsByUserIdAndMonth(int userId, String monthYear) throws SQLException {
        List<Budget> budgets = new ArrayList<>();
        String sql = "SELECT * FROM budgets WHERE user_id = ? AND month_year = ? ORDER BY category";
        try (Connection conn = DatabaseUtil.getConnection();
             PreparedStatement stmt = conn.prepareStatement(sql)) {
            stmt.setInt(1, userId);
            stmt.setString(2, monthYear);
            ResultSet rs = stmt.executeQuery();
            while (rs.next()) {
                budgets.add(extractBudgetFromResultSet(rs));
            }
        }
        return budgets;
    }

    // 获取用户指定分类和月份的预算
    public Budget getBudgetByUserCategoryAndMonth(int userId, String category, String monthYear) throws SQLException {
        String sql = "SELECT * FROM budgets WHERE user_id = ? AND category = ? AND month_year = ?";
        try (Connection conn = DatabaseUtil.getConnection();
             PreparedStatement stmt = conn.prepareStatement(sql)) {
            stmt.setInt(1, userId);
            stmt.setString(2, category);
            stmt.setString(3, monthYear);
            ResultSet rs = stmt.executeQuery();
            if (rs.next()) {
                return extractBudgetFromResultSet(rs);
            }
        }
        return null;
    }

    // 添加预算
    public boolean addBudget(Budget budget) throws SQLException {
        String sql = "INSERT INTO budgets (user_id, category, amount, month_year) VALUES (?, ?, ?, ?)";
        try (Connection conn = DatabaseUtil.getConnection();
             PreparedStatement stmt = conn.prepareStatement(sql)) {
            stmt.setInt(1, budget.getUserId());
            stmt.setString(2, budget.getCategory());
            stmt.setBigDecimal(3, budget.getAmount());
            stmt.setString(4, budget.getMonthYear());
            return stmt.executeUpdate() > 0;
        }
    }

    // 更新预算
    public boolean updateBudget(Budget budget) throws SQLException {
        String sql = "UPDATE budgets SET amount = ? WHERE budget_id = ?";
        try (Connection conn = DatabaseUtil.getConnection();
             PreparedStatement stmt = conn.prepareStatement(sql)) {
            stmt.setBigDecimal(1, budget.getAmount());
            stmt.setInt(2, budget.getBudgetId());
            return stmt.executeUpdate() > 0;
        }
    }

    // 删除预算
    public boolean deleteBudget(int budgetId) throws SQLException {
        String sql = "DELETE FROM budgets WHERE budget_id = ?";
        try (Connection conn = DatabaseUtil.getConnection();
             PreparedStatement stmt = conn.prepareStatement(sql)) {
            stmt.setInt(1, budgetId);
            return stmt.executeUpdate() > 0;
        }
    }

    // 获取用户指定月份的总预算
    public BigDecimal getTotalBudgetByUserIdAndMonth(int userId, String monthYear) throws SQLException {
        String sql = "SELECT SUM(amount) FROM budgets WHERE user_id = ? AND month_year = ?";
        try (Connection conn = DatabaseUtil.getConnection();
             PreparedStatement stmt = conn.prepareStatement(sql)) {
            stmt.setInt(1, userId);
            stmt.setString(2, monthYear);
            ResultSet rs = stmt.executeQuery();
            if (rs.next()) {
                BigDecimal total = rs.getBigDecimal(1);
                return total != null ? total : BigDecimal.ZERO;
            }
        }
        return BigDecimal.ZERO;
    }

    // 检查预算是否存在
    public boolean isBudgetExists(int userId, String category, String monthYear) throws SQLException {
        String sql = "SELECT COUNT(*) FROM budgets WHERE user_id = ? AND category = ? AND month_year = ?";
        try (Connection conn = DatabaseUtil.getConnection();
             PreparedStatement stmt = conn.prepareStatement(sql)) {
            stmt.setInt(1, userId);
            stmt.setString(2, category);
            stmt.setString(3, monthYear);
            ResultSet rs = stmt.executeQuery();
            if (rs.next()) {
                return rs.getInt(1) > 0;
            }
        }
        return false;
    }

    // 从结果集提取预算对象
    private Budget extractBudgetFromResultSet(ResultSet rs) throws SQLException {
        Budget budget = new Budget();
        budget.setBudgetId(rs.getInt("budget_id"));
        budget.setUserId(rs.getInt("user_id"));
        budget.setCategory(rs.getString("category"));
        budget.setAmount(rs.getBigDecimal("amount"));
        budget.setMonthYear(rs.getString("month_year"));
        budget.setCreatedAt(rs.getTimestamp("created_at").toLocalDateTime());
        budget.setUpdatedAt(rs.getTimestamp("updated_at").toLocalDateTime());
        return budget;
    }
}